
import numpy as np

import rclpy
from rclpy.node import Node
from std_msgs.msg import String


#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from geometry_msgs.msg import Twist  # 控制指令消息类型
from sensor_msgs.msg import Imu  # 示例传感器消息类型
import serial
import threading
import struct
import time

class ChassisDriver(Node):
    def __init__(self):
        super().__init__('chassis_driver')
        
        # 串口初始化参数
        self.declare_parameter('port', '/dev/ttyUSB0')  # 默认串口设备
        self.declare_parameter('baudrate', 115200)      # 默认波特率
        self.declare_parameter('cmd_topic', 'cmd_vel')  # 控制指令订阅话题
        
        # 初始化串口连接
        self._init_serial_connection()
        
        # 创建控制指令订阅器
        self.cmd_sub = self.create_subscription(
            Twist,
            self.get_parameter('cmd_topic').value,
            self.cmd_callback,
            10
        )
        
        # 创建数据发布器（示例：发布IMU数据）
        self.imu_pub = self.create_publisher(Imu, 'chassis/imu', 10)
        
        # 启动串口读取线程
        self.serial_thread = threading.Thread(target=self._serial_read_loop)
        self.serial_thread.daemon = True
        self.serial_thread.start()
        
        self.get_logger().info("底盘驱动节点已启动")

    def _init_serial_connection(self):
        """初始化串口连接"""
        try:
            self.ser = serial.Serial(
                port=self.get_parameter('port').value,
                baudrate=self.get_parameter('baudrate').value,
                timeout=0.1  # 非阻塞读取
            )
            self.get_logger().info(f"串口连接成功: {self.ser.name}")
        except Exception as e:
            self.get_logger().error(f"串口连接失败: {str(e)}")
            raise

    def cmd_callback(self, msg):
        """处理控制指令回调"""
        # 示例协议：vx,vy,vz (单位m/s)
        cmd_data = struct.pack('<fff', msg.linear.x, msg.linear.y, msg.angular.z)
        self._send_serial_command(b'CMD', cmd_data)
        self.get_logger().debug(f"发送指令: vx={msg.linear.x:.2f}, vz={msg.angular.z:.2f}")

    def _send_serial_command(self, header, data):
        """通过串口发送数据（带校验）"""
        try:
            # 协议帧格式: [0xAA][header][len][data][checksum]
            frame = bytearray()
            frame.append(0xAA)  # 帧头
            frame.extend(header)
            frame.append(len(data))
            frame.extend(data)
            
            # 计算校验和（简单异或校验）
            checksum = 0
            for byte in frame[1:]:
                checksum ^= byte
            frame.append(checksum)
            
            self.ser.write(frame)
        except Exception as e:
            self.get_logger().error(f"指令发送失败: {str(e)}")

    def _serial_read_loop(self):
        """串口数据读取线程"""
        buffer = bytearray()
        while rclpy.ok() and self.ser.is_open:
            try:
                data = self.ser.read(self.ser.in_waiting or 1)
                if data:
                    buffer.extend(data)
                    self._process_serial_data(buffer)
            except Exception as e:
                self.get_logger().error(f"串口读取错误: {str(e)}")
                time.sleep(0.1)

    def _process_serial_data(self, buffer):
        """处理原始串口数据"""
        # 示例协议帧：0xAA + 数据类型(3字节) + 长度(1字节) + 数据(N字节) + 校验(1字节)
        while len(buffer) >= 6:  # 最小帧长度检查
            start_idx = buffer.find(b'\xAA')
            if start_idx == -1:
                del buffer[:]
                return
            
            if len(buffer) < start_idx + 6:
                return  # 等待更多数据
                
            header = buffer[start_idx+1:start_idx+4]
            data_len = buffer[start_idx+4]
            
            # 检查完整帧是否接收完成
            if len(buffer) < start_idx + 6 + data_len:
                return
                
            # 提取数据帧
            frame = buffer[start_idx:start_idx+6+data_len]
            del buffer[:start_idx+6+data_len]  # 移除已处理数据
            
            # 校验检查
            checksum = 0
            for byte in frame[1:-1]:
                checksum ^= byte
                
            if checksum != frame[-1]:
                self.get_logger().warn("校验失败，丢弃帧")
                continue
                
            # 处理有效数据
            self._parse_data_frame(header, frame[5:5+data_len])

    def _parse_data_frame(self, header, data):
        """解析有效数据帧"""
        try:
            if header == b'IMU':
                # 示例：解析IMU数据 (ax,ay,az,wx,wy,wz)
                imu_data = struct.unpack('<ffffff', data)
                self._publish_imu_data(imu_data)
            # 添加其他数据类型解析...
        except Exception as e:
            self.get_logger().error(f"数据解析错误: {str(e)}")

    def _publish_imu_data(self, data):
        """发布IMU数据"""
        msg = Imu()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.header.frame_id = 'base_imu'
        msg.linear_acceleration.x = data
        msg.linear_acceleration.y = data
        msg.linear_acceleration.z = data
        msg.angular_velocity.x = data
        msg.angular_velocity.y = data
        msg.angular_velocity.z = data
        self.imu_pub.publish(msg)
        self.get_logger().debug("发布IMU数据")

    def destroy_node(self):
        """节点销毁时关闭串口"""
        if hasattr(self, 'ser') and self.ser.is_open:
            self.ser.close()
        super().destroy_node()

def main(args=None):
    rclpy.init(args=args)
    driver = ChassisDriver()
    try:
        rclpy.spin(driver)
    except KeyboardInterrupt:
        pass
    finally:
        driver.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
